import streamlit as st
import requests

# 设置页面配置
st.set_page_config(layout="wide")

# 设置页面样式
st.markdown("""
    <style>
        /* 设置对话区域的滚动样式 */
        .conversation-container {
            max-height: 500px; /* 设置最大高度 */
            overflow-y: auto;
            padding-right: 1rem;
        }
        
        /* 调整垂直间距 */
        [data-testid="stVerticalBlock"] {
            gap: 0rem;
        }
        
        /* 隐藏滚动条但保持功能 */
        .conversation-container::-webkit-scrollbar {
            width: 8px;
        }
        
        .conversation-container::-webkit-scrollbar-track {
            background: #f1f1f1;
        }
        
        .conversation-container::-webkit-scrollbar-thumb {
            background: #888;
            border-radius: 4px;
        }
        
        .conversation-container::-webkit-scrollbar-thumb:hover {
            background: #555;
        }
    </style>
""", unsafe_allow_html=True)

# 设置页面标题
st.title("RAG 智能问答系统")

# 创建左右两列布局
left_col, right_col = st.columns([1, 2])

# 左侧列 - 文档上传区域
with left_col:
    st.header("知识库文档上传")
    st.markdown("上传文档到知识库")
    
    uploaded_file = st.file_uploader(
        "支持 PDF、TXT、DOCX 格式",
        type=['pdf', 'txt', 'docx'],
        help="单个文件大小限制为200MB"
    )
    
    if uploaded_file:
        files = {"file": uploaded_file}
        with st.spinner("正在上传文件..."):
            response = requests.post("http://localhost:8000/upload", files=files)
            if response.status_code == 200:
                st.success("文档上传成功！")
            else:
                st.error("文档上传失败！")

# 右侧列 - 对话区域
with right_col:
    st.header("智能问答")
    
    # 初始化聊天历史
    if "messages" not in st.session_state:
        st.session_state.messages = []

    # 创建可滚动的对话容器
    with st.container():
        st.markdown('<div class="conversation-container">', unsafe_allow_html=True)
        
        # 显示聊天历史
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])
        
        st.markdown('</div>', unsafe_allow_html=True)

# 用户输入框
prompt = st.chat_input("请输入您的问题")

# 处理用户输入
if prompt:
    # 添加用户问题到聊天历史
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    # 在右侧显示用户消息和获取回答
    with right_col:
        with st.container():
            with st.chat_message("user"):
                st.markdown(prompt)
        
            # 显示助手回答
            with st.chat_message("assistant"):
                with st.spinner("思考中..."):
                    try:
                        response = requests.post(
                            "http://localhost:8000/chat",
                            json={"question": prompt},
                            timeout=30
                        )
                        
                        if response.status_code == 200:
                            answer = response.json()["answer"]
                            st.markdown(answer)
                            st.session_state.messages.append({"role": "assistant", "content": answer})
                        else:
                            st.error(f"获取回答失败！状态码：{response.status_code}")
                    except requests.exceptions.RequestException as e:
                        st.error(f"连接服务器失败：{str(e)}")

    # 使用 JavaScript 将滚动条保持在底部
    st.markdown("""
        <script>
            // 等待页面加载完成
            window.onload = function() {
                // 获取对话容器
                const container = document.querySelector('.conversation-container');
                // 将滚动条滚动到底部
                container.scrollTop = container.scrollHeight;
            };
        </script>
    """, unsafe_allow_html=True)

# 添加页面底部说明
st.markdown("---")
st.markdown("💡 提示：您可以通过上传文档来扩充知识库，提高回答的准确性。")
